# -*- coding: utf-8 -*-
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import minimize
import seaborn as sns

# 设置matplotlib支持中文显示
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams["axes.unicode_minus"] = False  # 正确显示负号


# 1. 数据准备
def load_preprocess_data():
    stocks = {
        '0016': '新鸿基地产',
        '0011': '恒生银行',
        '0001': '长和'
    }
    dfs = []

    for code, name in stocks.items():
        try:
            df = pd.read_csv(f'D:/MathD/{name}{code}历史数据.csv',
                             parse_dates=['日期'],
                             thousands=',')
            df['代码'] = code
            df['名称'] = name

            # 改进的交易量处理
            def convert_volume(v):
                if isinstance(v, str):
                    if 'M' in v:
                        return float(v.replace('M', '')) * 1e6
                    elif 'K' in v:
                        return float(v.replace('K', '')) * 1e3
                    else:
                        return float(v)
                return v

            df['交易量'] = df['交易量'].apply(convert_volume)
            dfs.append(df)
        except Exception as e:
            print(f"处理{name}{code}数据时出错:", str(e))

    if not dfs:
        raise ValueError("未能加载任何股票数据，请检查文件路径和格式")

    full_df = pd.concat(dfs)
    full_df = full_df.sort_values(['代码', '日期']).reset_index(drop=True)
    return full_df


# 2. 收益率计算函数
def calculate_returns(df, scheme=1):
    """计算两种方案下的日收益率"""
    if scheme == 1:
        df['买入价'] = df['低']
        df['卖出价'] = df['高']
    else:
        df['买入价'] = df['低'] + 0.25 * (df['高'] - df['低'])
        df['卖出价'] = df['低'] + 0.75 * (df['高'] - df['低'])

    # 考虑T+1和交易费用
    df['收益率'] = (df.groupby('代码')['卖出价'].shift(-1) - df['买入价']) / df['买入价'] - 0.0003 * 2
    return df.dropna(subset=['收益率'])


# 3. 组合优化核心函数
class PortfolioOptimizer:
    def __init__(self, returns_df):
        self.returns = returns_df.pivot(index='日期', columns='代码', values='收益率')
        self.cov = self.returns.cov()
        self.mean_returns = self.returns.mean()
        self.n_assets = len(self.mean_returns)

    def portfolio_performance(self, weights):
        """计算组合的预期收益和风险"""
        port_return = np.sum(self.mean_returns * weights)
        port_vol = np.sqrt(np.dot(weights.T, np.dot(self.cov, weights)))
        sharpe = (port_return - 0.02) / port_vol  # 假设无风险利率2%
        return port_return, port_vol, sharpe

    def minimize_risk(self, target_return):
        """给定目标收益下的最小风险"""

        def objective(weights):
            return np.dot(weights.T, np.dot(self.cov, weights))

        constraints = [
            {'type': 'eq', 'fun': lambda x: np.sum(x) - 1},
            {'type': 'eq', 'fun': lambda x: np.sum(x * self.mean_returns) - target_return}
        ]
        bounds = tuple((0, 1) for _ in range(self.n_assets))
        init_guess = [1 / self.n_assets] * self.n_assets
        result = minimize(objective, init_guess,
                          method='SLSQP',
                          bounds=bounds,
                          constraints=constraints)
        return result

    def maximize_sharpe(self):
        """最大化夏普比率"""

        def objective(weights):
            ret, vol, _ = self.portfolio_performance(weights)
            return -ret / vol  # 最小化负夏普

        constraints = {'type': 'eq', 'fun': lambda x: np.sum(x) - 1}
        bounds = tuple((0, 1) for _ in range(self.n_assets))
        init_guess = [1 / self.n_assets] * self.n_assets
        result = minimize(objective, init_guess,
                          method='SLSQP',
                          bounds=bounds,
                          constraints=constraints)
        return result


# 4. 可视化函数
def plot_results(optimizer, scheme_name):
    """绘制有效前沿和最优组合"""
    # 生成有效前沿
    target_returns = np.linspace(optimizer.mean_returns.min(),
                                 optimizer.mean_returns.max(), 50)
    efficient_portfolios = []

    for ret in target_returns:
        res = optimizer.minimize_risk(ret)
        if res.success:
            _, vol, _ = optimizer.portfolio_performance(res.x)
            efficient_portfolios.append((ret, vol))

    if not efficient_portfolios:
        print("无法生成有效前沿，数据可能存在问题")
        return None

    eff_ret, eff_vol = zip(*efficient_portfolios)

    # 最优夏普组合
    sharpe_res = optimizer.maximize_sharpe()
    if not sharpe_res.success:
        print("无法找到最优夏普组合")
        return None

    sharpe_ret, sharpe_vol, _ = optimizer.portfolio_performance(sharpe_res.x)

    plt.figure(figsize=(12, 6))
    plt.scatter(np.sqrt(np.diag(optimizer.cov)), optimizer.mean_returns,
                c='blue', marker='o', s=200, label='个股')
    plt.plot(eff_vol, eff_ret, 'r--', linewidth=2, label='有效前沿')
    plt.scatter(sharpe_vol, sharpe_ret, c='green', marker='*', s=300,
                label='最优夏普组合')

    # 标注股票名称
    for i, txt in enumerate(optimizer.returns.columns):
        plt.annotate(txt, (np.sqrt(optimizer.cov.iloc[i, i]),
                           optimizer.mean_returns[i]),
                     fontsize=12)

    plt.title(f'{scheme_name}投资组合优化结果', fontsize=14)
    plt.xlabel('波动率', fontsize=12)
    plt.ylabel('预期收益率', fontsize=12)
    plt.legend()
    plt.grid(True)
    plt.show()

    # 绘制权重饼图
    plt.figure(figsize=(6, 6))
    labels = [f'{code} ({name})' for code, name in zip(
        optimizer.returns.columns, ['新鸿基', '恒生', '长和'])]
    plt.pie(sharpe_res.x, labels=labels, autopct='%1.1f%%')
    plt.title('最优组合权重分配')
    plt.show()

    return sharpe_res.x


# 5. 主程序
def main():
    try:
        # 加载并预处理数据
        df = load_preprocess_data()

        # 检查数据是否加载成功
        if df.empty:
            print("数据加载失败，请检查文件路径和格式")
            return

        # 对两种方案分别处理
        for scheme in [1, 2]:
            scheme_name = '方案1(最高最低价)' if scheme == 1 else '方案2(分位点定价)'
            print(f"\n===== {scheme_name} =====")

            # 计算收益率
            returns_df = calculate_returns(df.copy(), scheme)

            # 检查收益率数据
            if returns_df.empty:
                print(f"{scheme_name}收益率计算失败，跳过此方案")
                continue

            # 初始化优化器
            optimizer = PortfolioOptimizer(returns_df)

            # 输出基本统计信息
            print("\n各股票年化收益率:")
            print((1 + optimizer.mean_returns) **252 - 1)

            print("\n协方差矩阵:")
            print(optimizer.cov)

            # 优化并可视化
            optimal_weights = plot_results(optimizer, scheme_name)

            if optimal_weights is not None:
                # 输出最优组合
                optimal_return, optimal_vol, optimal_sharpe = optimizer.portfolio_performance(optimal_weights)
                print(f"\n最优组合权重: {dict(zip(optimizer.returns.columns, optimal_weights))}")
                print(f"年化收益率: {(1 + optimal_return)** 252 - 1:.2%}")
                print(f"年化波动率: {optimal_vol * np.sqrt(252):.2%}")
                print(f"夏普比率: {optimal_sharpe:.2f}")
    except Exception as e:
        print("程序运行出错:", str(e))


if __name__ == '__main__':
    main()
